# main.py
import os
from db import connect_to_database
from agent import create_agent
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

# 定义 API 请求体
class QueryRequest(BaseModel):
    query: str

# 定义 API 响应体
class QueryResponse(BaseModel):
    response: str

# 初始化数据库连接
db_connection = connect_to_database(
    host='XXX.XXX.XXX.XXX',
    database='dbname',
    user='username',
    password='password'
)

if not db_connection:
    raise Exception("无法连接到数据库")

# 创建Agent


# 这里没有启用统一的key设置，可以自行修改传入参数，避免多处进行配置。
openai_api_key="default"

# 初始化FastAPI应用
app = FastAPI()

@app.post("/query", response_model=QueryResponse)
def query_database(request: QueryRequest):
    user_query = request.query
    # user_input="查找node有几种type，每种type有多少数据"
    print("user_query:"+user_query)
    try:
        result =  create_agent(openai_api_key, db_connection,user_query)
        print(result)
        return QueryResponse(response=result)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))